# query_vector.py
# python ./src/query/query_vector.py --query_text "请问在2025年的农林牧渔专题中，D区办结工单量相较于A区多出多少（若A区更多则请用负数表示）？" --debug
# 唯一TOP01： python ./src/query/query_vector.py --query_text "请问2023年3月交通运输专题工单反馈的工单量排名前4的事项有哪些？"
# TOP03：python ./src/query/query_vector.py --query_text "请问在2025年的农林牧渔专题中，D区办结工单量相较于A区多出多少（若A区更多则请用负数表示）？"
# python ./src/query/query_vector.py --query_text "请问2024年区县单位平均处理工单时长是多少？"

import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
import chromadb
import numpy as np
from numpy.linalg import norm
from query.hit_keywords import hit_keywords
import argparse
import json
from model.ai_loader import clientEmbedding
from config import project_config
from utils import debug_print


# 配置
CHROMA_DB_PATH = project_config.chroma_location
COLLECTION_NAME = project_config.collection_name

def cosine_similarity(a, b):
    """计算两个向量的余弦相似度"""
    return np.dot(a, b) / (norm(a) * norm(b))

def get_query_embedding(texts, model=project_config.embedding_model):

    try:
        response = clientEmbedding.embeddings.create(
            input=texts,
            model=model
        )
        # 提取嵌入向量
        embeddings = [data.embedding for data in response.data]
        return embeddings[0]
    except Exception as e:
        print(f"生成嵌入时出错: {e}")
        return None

def retrieve_similar_texts(query_text="helloworld", debug=None, top_k=5) : # , silent=False):
    """检索与 query_text 最相似的文本"""
    # 连接到持久化 Chroma 客户端
    client = chromadb.PersistentClient(path=CHROMA_DB_PATH)
    # 获取集合（无需 embedding function，因为我们直接使用已有 embedding）
    collection = client.get_collection(name=COLLECTION_NAME)
    # Step 1: 获取查询文本的 embedding（必须与 log 中维度一致）
    try:
        query_embedding = get_query_embedding(query_text)
        if len(query_embedding) != project_config.embedding_dim_size:
            raise ValueError(f"Embedding 维度错误：期望 {project_config.embedding_dim_size} {len(query_embedding)}")
    except Exception as e:
        print(f"生成查询 embedding 失败：{e}")
        return
    # Step 2: 使用 query_embedding 进行查询
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "distances", "embeddings"]  # 可选：返回距离和 embedding
    )
    # Step 3: 解析并打印结果
    debug_print(debug, f"\n🔍 查询文本: '{query_text}'\n")
    debug_print(debug, f"找到 {len(results['documents'][0])} 个最相似的记录：\n")
    # debug_print(debug, f"results: {results['documents'][0]} \n")
    # 注意：query 返回的是 list of list，因为支持批量查询
    docs = results['documents'][0]
    distances = results['distances'][0]  # 距离（越小越相似）
    ids = results['ids'][0]
    # Chroma 默认使用余弦距离（1 - cosine_similarity），所以转换回相似度
    similarities = [1 - d for d in distances]
    filtered_results = [
        (doc, sim, idx)
        for doc, sim, idx in zip(docs, similarities, ids)
        if sim >= 0
    ]
    debug_print(debug, f"找到 {len(filtered_results)} 个相似度 ≥ 0 的最相似记录：\n")
    for i, (doc_id, text, sim) in enumerate(zip(ids, docs, similarities)):
        # 将 text 按行分割，取前 3 行
        lines = text.splitlines()
        text_preview = '\n'.join(lines[:3])
        debug_print(debug, f"Rank {i+1}:")
        debug_print(debug, f"  ID: {doc_id}")
        debug_print(debug, f"  Similarity: {sim:.4f}")
        debug_print(debug, f"  Text: {text_preview}")
        debug_print(debug, "-" * 50)
        filtered_docs = [item[0] for item in filtered_results]  # 只取 doc 文本
    hitIndexs = hit_keywords(query_text, filtered_docs, debug)
    # 继续原有代码逻辑，处理 hitIndexs 的返回结果
    if hitIndexs is None or len(hitIndexs) == 0:
        debug_print(debug, "⚠️ 未命中任何关键词，无法判断优先级。")
        filtered_docs = results['documents'][0]
        hitIndexs = hit_keywords(query_text, filtered_docs, debug)
        # return filtered_docs 
    # 检查 hitIndexs 中是否存在唯一的最大值
    max_count = max(hitIndexs)
    max_count_occurrences = hitIndexs.count(max_count)
    if max_count_occurrences >= 1:
        # 多个最大值：保留前三个（即前三个命中最高），删除其余
        top_k_to_keep = 5
        # 找出所有具有最大值的索引，并取前 top_k_to_keep 个
        top_indices = [
            i for i, v in enumerate(hitIndexs) if v == max_count
        ][:top_k_to_keep]
        # 保留 filtered_docs 中对应的记录
        filtered_docs = [filtered_docs[i] for i in top_indices]
        return filtered_docs

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description="检索向量库中相似度最高的内容。")
    parser.add_argument('--query_text', type=str, required=True, help='需要检索的文本字符串。')
    parser.add_argument('--debug', action='store_true', help='启用调试模式，显示详细日志。')
    args = parser.parse_args()
    text_to_query = args.query_text

    # 将 silent 参数传递给函数
    filtered_docs = retrieve_similar_texts(text_to_query,args.debug, top_k=15) # , silent=args.silent)
    print(f"filtered_docs={filtered_docs}")
    # 提取所有 document，生成新的 list

    try:
        with open("tmp_filtered_docs.txt", "w", encoding="utf-8") as f:
            json.dump(filtered_docs, f, ensure_ascii=False, indent=4)
        print(f"\n✅ [DEBUG] 已将查询和结果写入 'tmp_filtered_docs.txt'")
    except Exception as e:
        print(f"❌ 写入文件失败: {e}")

# --- 程序入口 ---
if __name__ == '__main__':
    main()
